# -*- coding: utf-8 -*-
'''
Created on 01.08.2019

@author: yu03
'''
from PyUeye import file_name
import numpy as np
import os
import matplotlib.pyplot as plt
import cv2
from scipy import signal
from FFT_Interpolation import FFT_interpolation_boxcar, FFT_interpolation_2, FFT_interpolation_compare, FFT_cal
from scipy.optimize import curve_fit
from scipy import signal
import datetime
now = datetime.datetime.now()

def fit_func(x, a, b, c):
    return a*(x-b)**2 + c
# p= -0.0213125
p= -0.02135
def fit_func_phase(x, q):
    return p*x+q
def fit_func_line(x, p, q):
    return p*x+q

#############
# fs_cam = 575.8  ###USB2.0
# fs_cam = 1370.72 ###USB3.0
# fs_cam = 1860 ###USB3.0 numpy readout
fs_cam = 1000 ### Hardware Trigger
####################

Lamda = 633e-9
pix_size = 5.3e-6
V_x, V_y, V_z = 0, 0, 0

img_set = []
hor_phase_centers = []
hor_freq_set = []
hor_f_fit_set = []
hor_phi_fit_set = []

f = open(file_name, 'rb')







f_size = os.fstat(f.fileno()).st_size
frames = []
time_sequence = []
while f.tell() < f_size:
#         print(f.tell())
    frame = np.load(f, allow_pickle=True)
    frames.append(frame[0])
    time_sequence.append(frame[1])

# x_axis = np.linspace(0, len(time_sequence)-1, len(time_sequence))

frames = np.array(frames)
print(frames.shape)

if len(frames.shape) == 4:
    print('file_size: %.1f MB'%(f_size/1e6))
    print('%i frames'%len(frames))
    print('%i lines/frame'%len(frames[0]))
    print('%i pixels/line'%len(frames[0][0]))
    print('%i channels/pixel'%len(frames[0][0][0]))

elif len(frames.shape) == 3:
    print('file_size: %.1fMB'%(f_size/1e6))
    print('%i frames'%len(frames))
    print('1 line/frame')
    print('%i pixels/line'%len(frames[0]))
    print('%i channels/pixel'%len(frames[0][0]))

print('frame rate: %.2f'%fs_cam)
print('record time: %f s'%(len(frames)/fs_cam))
# print(len(frames[0][0]))
# frame = frames[0]
# line = frame[0]
# line = np.array([i[0]+256*i[1] for i in line])
# print(line)

num = 0
time_sequence = time_sequence[::]
for frame in frames[::10]:
    num += 1
    if len(frames.shape) == 3:
        line = frame
    elif len(frames.shape) == 4:
        line = frame[0]
#     line = frame[0]

    line = np.array([i[0]+256*i[1] for i in line])
    line = np.diff(line.astype('float32'))
    line = np.concatenate(([0], line))
    hor_center = line
     
    DC_num = 0
    hor_freq_estim, hor_phase_estim, hor_freqline, hor_sig_magnitude, hor_sig_phase,  hor_m_k_num, hor_X_m_k, hor_freq_for_phase = FFT_interpolation_2(hor_center, pix_size, 1e5, DC_num)
    hor_FFT_start = np.where(hor_sig_magnitude[DC_num:] > hor_X_m_k*0.4)[0][0]+DC_num
    hor_FFT_end = np.where(hor_sig_magnitude[DC_num:] > hor_X_m_k*0.4)[0][-1]+DC_num
    hor_fit_x = hor_freqline[hor_FFT_start:hor_FFT_end+1]
    hor_fit_y = hor_sig_magnitude[hor_FFT_start:hor_FFT_end+1]
    hor_fit_phase = hor_sig_phase[hor_FFT_start:hor_FFT_end+1]
    hor_params = curve_fit(fit_func, hor_fit_x, hor_fit_y)
    [hor_a, hor_b, hor_c] = hor_params[0]
    hor_f_fit = hor_b
    hor_f_fit_set.append(hor_f_fit)
#     hor_params = curve_fit(fit_func_phase, hor_fit_x, np.unwrap(hor_fit_phase))
#     [hor_q] = hor_params[0]
#     hor_phi_fit = p*hor_freqline[hor_m_k_num]+hor_q
#     hor_phi_fit_set.append(hor_phi_fit+2*np.pi-0.065)
    print(num, hor_m_k_num)
      
    hor_freq_set.append(hor_freq_estim)    
    hor_phase_centers.append(hor_phase_estim)
    
hor_f_fit_set = np.array(hor_f_fit_set)
hor_phase_centers = np.array(hor_phase_centers)
hor_angle = (V_x-Lamda*hor_f_fit_set/2)*1e6 ### urad
hor_length = np.unwrap(hor_phase_centers)/4/np.pi*Lamda*1e9 ### nm

plt.figure('line')
plt.plot(hor_center)

if 1: ### Plotting Noise_PSD & Square wave
    
    plt.figure(1)
    plt.gcf().set_size_inches(18,9)
    ax1 = plt.subplot(3,2,1)
    cut_range = [0, len(hor_angle)]
    x = (np.arange(cut_range[1]-cut_range[0])+cut_range[0])/fs_cam
    y = hor_angle
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Horizontal Tilting")
    plt.ylabel("Horizontal Angle (urad)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
        
    ax2 = plt.subplot(3,2,2)
    cut_range = [0, len(hor_length)]
    x = (np.arange(cut_range[1]-cut_range[0])+cut_range[0])/fs_cam
    y = hor_length
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Horizontal Phase")
    plt.ylabel("Horizontal Length (nm)]")
    plt.xlabel("Samples")
    plt.grid(which='both', axis='both')
      
    ax3 = plt.subplot(3,2,3)
#     cut_range = [2700, 3200]
    x = (np.arange(cut_range[1]-cut_range[0])+cut_range[0])/fs_cam
    y = (hor_angle-np.average(hor_angle))[cut_range[0]:cut_range[1]]
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Horizontal Tilting")
    plt.ylabel("Horizontal Angle (urad)")
    plt.xlabel("Samples")
    plt.grid(which='both', axis='both')
       
    ax4 = plt.subplot(3,2,4)
#     cut_range = [2700, 3200]
    x = (np.arange(cut_range[1]-cut_range[0])+cut_range[0])/fs_cam
    y = (hor_length-np.average(hor_length))[cut_range[0]:cut_range[1]]
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Horizontal Length")
    plt.ylabel("Horizontal Length (nm)")
    plt.xlabel("Samples")
    plt.grid(which='both', axis='both')

    ax5 = plt.subplot(3,2,5)
    f_line, PS = signal.welch(hor_angle, fs_cam, nperseg=len(hor_angle))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Horizontal Tilting Noise_PSD")
    plt.ylabel("PSD (urad/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-6, 10) # for Ramp/20s noise
#     plt.ylim(1e-4, 0.1) # for noise
    
    ax6 = plt.subplot(3,2,6)
    f_line, PS = signal.welch(hor_length, fs_cam, nperseg=len(hor_length))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Horizontal Length Noise_PSD")
    plt.ylabel("PSD (nm/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-6, 10) # for Ramp/20s noise
#     plt.ylim(1e-6, 2e-2) # for square wave
#     plt.ylim(1e-6, 1e-3)# for noise
    
    plt.tight_layout()
    plt.show()
    
    
if 0: ### Plotting Nonlinearity
    
    ############### linearity range
    start_cutting, stop_cutting = 0, len(hor_phase_centers)
    start_cutting, stop_cutting = start_cutting/fs_cam, stop_cutting/fs_cam
    start_cutting, stop_cutting = 0, 0.8
    zoom_time = start_cutting, stop_cutting
    
#     start_cutting, stop_cutting = 1, 1.5 ### in second   
#     zoom_time = [1.2, 1.4]
    
    start_cutting, stop_cutting = int(start_cutting*fs_cam), int(stop_cutting*fs_cam)
    selected_range = np.linspace(start_cutting, stop_cutting, stop_cutting-start_cutting)
    selected_time = (np.arange(stop_cutting-start_cutting)+start_cutting)/fs_cam    
    zoom_range = [zoom_time[0]*fs_cam, zoom_time[1]*fs_cam]
    
    
    ############### Plotting
     
    plt.figure(1)
    plt.gcf().set_size_inches(18,9)
    ax1 = plt.subplot(3,2,1)
    range = [0, len(hor_angle)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = hor_angle
    plt.plot(x, y, color='blue', marker=' ')
    measured_angle = hor_angle[start_cutting:stop_cutting]
    [p, q] = np.polyfit(np.linspace(0, len(measured_angle)-1, len(measured_angle)), measured_angle, 1)
    linear_fit = np.linspace(0, len(measured_angle)-1, len(measured_angle))*p + q
    hor_angle_nonlinearity = (measured_angle - linear_fit)
    plt.plot(selected_time, linear_fit, color='red')
    plt.title("Horizontal Tilting")
    plt.ylabel("Horizontal Angle (urad)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
        
    ax2 = plt.subplot(3,2,2)
    range = [0, len(hor_length)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = hor_length
    plt.plot(x, y, color='blue', marker=' ')
    
    measured_length = hor_length[start_cutting:stop_cutting]
    [p, q] = np.polyfit(np.linspace(0, len(measured_length)-1, len(measured_length)), measured_length, 1)
    linear_fit = np.linspace(0, len(measured_length)-1, len(measured_length))*p + q
    hor_length_nonlinearity = (measured_length - linear_fit)
    plt.plot(selected_time, linear_fit, color='red')
    plt.title("Horizontal Length")
    plt.ylabel("Horizontal Length (nm)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
      
    ax3 = plt.subplot(3,2,3)
    range = [start_cutting, stop_cutting]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = hor_angle_nonlinearity
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Horizontal Tilting")
    plt.ylabel("Horizontal Angle (urad)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
       
    ax4 = plt.subplot(3,2,4)
    range = [start_cutting, stop_cutting]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = hor_length_nonlinearity
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Horizontal Length")
    plt.ylabel("Horizontal Length (nm)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
      
    ax5 = plt.subplot(3,2,5)
    f_line, PS = signal.welch(hor_angle_nonlinearity, fs_cam, nperseg=len(hor_angle_nonlinearity))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Horizontal Tilting Noise_PSD")
    plt.ylabel("PSD (urad/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-4, 10) # for Ramp/20s noise
#     plt.ylim(1e-4, 0.1) # for noise
    
    ax6 = plt.subplot(3,2,6)
#     measured_length = (hor_length+ver_length)/2
#     [p, q] = np.polyfit(np.linspace(0, len(measured_length)-1, len(measured_length)), measured_length, 1)
#     linear_fit = np.linspace(0, len(measured_length)-1, len(measured_length))*p + q
#     length_nonlinearity = (measured_length - linear_fit)
#     plt.plot(x, length_nonlinearity, color='blue', marker=' ')
    f_line, PS = signal.welch(hor_length_nonlinearity, fs_cam, nperseg=len(hor_length_nonlinearity))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Horizontal Length Noise_PSD")
    plt.ylabel("PSD (nm/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-4, 1e1) # for Ramp/20s noise
    
    plt.tight_layout()
    plt.show()

    
if 0: ### Plotting
    
    ############### linearity range
    start_cutting, stop_cutting = 0, len(hor_phase_centers)
#     start_cutting, stop_cutting = 20, 3000
    
    selected_range = np.linspace(start_cutting, stop_cutting, stop_cutting-start_cutting)
     
    ############### Plotting
    plt.figure(1)
    plt.subplot(2,1,1)
    plt.plot(np.diff(time_sequence), 'b', marker='o')
    # plt.plot(linear_fit, 'r', marker=' ')
    plt.grid(which='both', axis='both')
    plt.subplot(2,1,2)
    plt.plot(time_sequence)
    plt.grid(which='both', axis='both')
     
    plt.figure(2)
    plt.plot(line)
    plt.grid(which='both', axis='both')
      
    scale_timestamp = 2000
     
    plt.figure(3)
    plt.gcf().set_size_inches(18,9)
    ax1 = plt.subplot(4,2,1)
    plt.plot(hor_f_fit_set, color='blue', marker=' ')
#     plt.plot(scale_timestamp*(np.diff(time_sequence)-np.average(np.diff(time_sequence)))+np.average(hor_f_fit_set), color='black')
    measured_freq = hor_f_fit_set[start_cutting:stop_cutting]
    # fit_params = curve_fit(fit_func_line, np.linspace(0, len(measured_freq)-1, len(measured_freq)), measured_freq)
    # [p, q] = fit_params[0]
    [p, q] = np.polyfit(np.linspace(0, len(measured_freq)-1, len(measured_freq)), measured_freq, 1)
    linear_fit = np.linspace(0, len(measured_freq)-1, len(measured_freq))*p + q
    hor_freq_nonlinearity = (measured_freq - linear_fit)
    plt.plot(selected_range, linear_fit, color='red')
    plt.title("Horizontal FFT")
    plt.ylabel("Horizontal Freq Estimation (/m)")
    plt.xlabel("Samples")
    plt.grid(which='major', axis='both')
    freq_range = plt.gca().get_ylim()
    angle_range = ((V_x-Lamda*freq_range[0]/2)*1e6, (V_x-Lamda*freq_range[1]/2)*1e6)
    ax1_angle = ax1.twinx()
    plt.ylim(angle_range)
    plt.ylabel('Horizontal Tilting (urad)')
        
    ax2 = plt.subplot(4,2,2)
    hor_phase_unwrap = np.unwrap(hor_phase_centers)
    plt.plot(hor_phase_unwrap, color='blue', marker=' ')
#     plt.plot(scale_timestamp/100*(np.diff(time_sequence)-np.average(np.diff(time_sequence)))+np.average(hor_phase_unwrap), color='black')
    # plt.plot(hor_phi_fit_set, color='red', marker=' ')
    measured_phase = hor_phase_unwrap[start_cutting:stop_cutting]
    #     fit_params = curve_fit(fit_func_line, np.linspace(0, len(measured_phase)-1, len(measured_phase)), np.unwrap(measured_phase))
    # [p, q] = fit_params[0]
    [p, q] = np.polyfit(np.linspace(0, len(measured_phase)-1, len(measured_phase)), np.unwrap(measured_phase), 1)
    linear_fit = np.linspace(0, len(measured_phase)-1, len(measured_phase))*p + q
    hor_phase_nonlinearity = (measured_phase - linear_fit)
    plt.plot(selected_range, linear_fit, color='red')
    plt.title("Horizontal Phase")
    plt.ylabel("Horizontal Phase (rad)")
    plt.xlabel("Samples")
    plt.grid(which='major', axis='both')
    phase_range = plt.gca().get_ylim()
    length_range = (phase_range[0]/4 /np.pi * Lamda*1e9, phase_range[1]/4 /np.pi * Lamda*1e9)
    ax2_length = ax2.twinx()
    plt.ylim(length_range)
    plt.ylabel('Horizontal Length (nm)')
      
    ax3 = plt.subplot(4,2,3)
    plt.plot(selected_range, hor_freq_nonlinearity, color='blue')
#     plt.plot(selected_range[1:], scale_timestamp*(np.diff(time_sequence[start_cutting:stop_cutting])-np.average(np.diff(time_sequence[start_cutting:stop_cutting])))+np.average(hor_freq_nonlinearity[1:]), color='black')
    plt.title("Horizontal Tilting Noise_PSD")
    plt.ylabel("Horizontal Freq Estimation (/m)")
    plt.xlabel("Samples")
    plt.grid(which='major', axis='both')
    freq_range = plt.gca().get_ylim()
    angle_range = ((V_x-Lamda*freq_range[0]/2)*1e6, (V_x-Lamda*freq_range[1]/2)*1e6)
    ax3_angle = ax3.twinx()
    plt.ylim(angle_range)
    plt.ylabel('Horizontal Tilting (urad)')
       
    ax4 = plt.subplot(4,2,4)
#     plt.plot(selected_range[1:], scale_timestamp/100*(np.diff(time_sequence[start_cutting:stop_cutting])-np.average(np.diff(time_sequence[start_cutting:stop_cutting])))+np.average(hor_phase_nonlinearity[1:]), color='black')
    plt.plot(selected_range, hor_phase_nonlinearity, color='blue')
#     plt.plot(selected_range, np.diff(hor_phase_unwrap)[start_cutting:stop_cutting], color='blue')
    plt.title("Horizontal Phase Noise_PSD")
    plt.ylabel("Horizontal Phase (rad)")
    plt.xlabel("Samples")
    plt.grid(which='major', axis='both')
    phase_range = plt.gca().get_ylim()
    length_range = (phase_range[0]/4 /np.pi * Lamda*1e9, phase_range[1]/4 /np.pi * Lamda*1e9)
    ax4_length = ax4.twinx()
    plt.ylim(length_range)
    plt.ylabel('Horizontal Length (nm)')
      
    ax5 = plt.subplot(4,2,5)
    plt.plot(selected_range, hor_freq_nonlinearity, color='blue')
#     plt.plot(selected_range[1:], scale_timestamp*(np.diff(time_sequence[start_cutting:stop_cutting])-np.average(np.diff(time_sequence[start_cutting:stop_cutting])))+np.average(hor_freq_nonlinearity), color='black')
    plt.title("Horizontal Tilting Noise_PSD")
    plt.ylabel("Horizontal Freq Estimation (/m)")
    plt.xlabel("Samples")
    plt.grid(which='major', axis='both')
    freq_range = plt.gca().get_ylim()
    angle_range = ((V_x-Lamda*freq_range[0]/2)*1e6, (V_x-Lamda*freq_range[1]/2)*1e6)
    ax5_angle = ax5.twinx()
    plt.ylim(angle_range)
    plt.ylabel('Horizontal Tilting (urad)')
      
    ax6 = plt.subplot(4,2,6)
    plt.plot(selected_range, hor_phase_nonlinearity, color='blue')
#     plt.plot(selected_range[1:], scale_timestamp/100*(np.diff(time_sequence[start_cutting:stop_cutting])-np.average(np.diff(time_sequence[start_cutting:stop_cutting])))+np.average(hor_phase_nonlinearity[1:]), color='black')
 
    plt.title("Horizontal Phase Noise_PSD")
    plt.ylabel("Horizontal Phase (rad)")
    plt.xlabel("Samples")
    plt.grid(which='major', axis='both')
    phase_range = plt.gca().get_ylim()
    length_range = (phase_range[0]/4 /np.pi * Lamda*1e9, phase_range[1]/4 /np.pi * Lamda*1e9)
    ax6_length = ax6.twinx()
    plt.ylim(length_range)
    plt.ylabel('Horizontal Length (nm)')
      
    # plt.tight_layout(pad=0.4, w_pad=0.5, h_pad=1.0)
      
    freqline, sig_FFT, sig_magnitude, sig_phase = FFT_cal(hor_f_fit_set, 1/fs_cam)
    # plt.figure('Noise_PSD Spectrum')
    ax1 = plt.subplot(4,2,7)
    f_line, PS = signal.welch(hor_freq_nonlinearity, fs_cam, nperseg=len(hor_freq_nonlinearity))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Freq. Noise_PSD FFT")
    plt.ylabel("Freq. Noise_PSD (/m/sqrt(Hz)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-4, 2) # for Ramp/20s noise
#     plt.ylim(1e-4, 0.2) # for noise 
    freq_range = plt.gca().get_ylim()
    angle_range = ((-V_x+Lamda*freq_range[0]/2)*1e6, (-V_x+Lamda*freq_range[1]/2)*1e6)
    ax1_angle = ax1.twinx()
    plt.ylim(angle_range)
    plt.yscale('Log')
    plt.ylabel('Tilting Amplitude (urad/sqrt(Hz))')
    ax2 = plt.subplot(4,2,8)
    f_line, PS = signal.welch(hor_phase_nonlinearity, fs_cam, nperseg=len(hor_phase_nonlinearity))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Phase Noise_PSD FFT")
    plt.ylabel("Phase Noise_PSD(rad/sqrt(Hz))")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-6, 1e-2) # for Ramp/ 20s noise
#     plt.ylim(1e-6, 2e-2) # for square wave
#     plt.ylim(1e-6, 1e-3)# for noise
    phase_range = plt.gca().get_ylim()
    length_range = (phase_range[0]/4 /np.pi * Lamda*1e9, phase_range[1]/4 /np.pi * Lamda*1e9)
    ax2_length = ax2.twinx()
    plt.ylim(length_range)
    plt.yscale('log')
    plt.ylabel('Length Amplitude (nm/sqrt(Hz))')
    plt.tight_layout()
    plt.show()